#include <opencv2/opencv.hpp>

#include "final_exam/kalman_predictor.h"

void Predictor::setDeltaT(double dt) { this->delta_t = dt; }

void Predictor::initKalmanFilter(const cv::Point3f &pt) {
  // 状态向量 [x, y, z, vx, vy, vz]
  this->state = cv::Mat::zeros(6, 1, CV_32F);
  this->state.at<float>(0) = pt.x;
  this->state.at<float>(1) = pt.y;
  this->state.at<float>(2) = pt.z;
  this->state.at<float>(3) = 0;
  this->state.at<float>(4) = 0;
  this->state.at<float>(5) = 0;

  // 测量向量 [z_x, z_y, z_z]
  this->meas = cv::Mat::zeros(3, 1, CV_32F);

  // 初始化 Kalman 滤波器
  this->kf.init(6, 3, 0, CV_32F);

  // 状态转移矩阵
  this->kf.transitionMatrix =
      (cv::Mat_<float>(6, 6) << 1, 0, 0, delta_t, 0, 0, 0, 1, 0, 0, delta_t, 0,
       0, 0, 1, 0, 0, delta_t, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0,
       0, 1);

  // 测量矩阵
  this->kf.measurementMatrix = cv::Mat::zeros(3, 6, CV_32F);
  this->kf.measurementMatrix.at<float>(0, 0) = 1.0f;
  this->kf.measurementMatrix.at<float>(1, 1) = 1.0f;
  this->kf.measurementMatrix.at<float>(2, 2) = 1.0f;

  // 初始化误差协方差矩阵
  cv::setIdentity(this->kf.errorCovPost, cv::Scalar::all(1));

  // 测量噪声协方差矩阵
  cv::setIdentity(this->kf.measurementNoiseCov, cv::Scalar::all(1e-2));

  // 过程噪声协方差矩阵
  cv::setIdentity(this->kf.processNoiseCov, cv::Scalar::all(1e-4));
  this->kf.statePost = this->state.clone();
  this->init = true;
  this->missing_count = 0;
}

cv::Point3f Predictor::predict() {
  if (!this->init) {
    throw std::runtime_error("Kalman Filter not initialized!");
  }

  // 预测下一位置
  cv::Mat prediction = this->kf.predict();
  cv::Point3f predictPt;
  predictPt.x = prediction.at<float>(0);
  predictPt.y = prediction.at<float>(1);
  predictPt.z = prediction.at<float>(2);

  // 增加丢失计数
  this->missing_count++;
  this->state = prediction.clone();
  return predictPt;
}

cv::Point3f Predictor::correct(const cv::Point3f &pt) {
  if (!this->init) {
    throw std::runtime_error("Kalman Filter not initialized!");
  }
  // 更新测量值
  this->meas.at<float>(0) = pt.x;
  this->meas.at<float>(1) = pt.y;
  this->meas.at<float>(2) = pt.z;
  // 用新的测量值更新卡尔曼滤波器
  cv::Mat estimated = this->kf.correct(this->meas);
  cv::Point3f statePt;
  statePt.x = estimated.at<float>(0);
  statePt.y = estimated.at<float>(1);
  statePt.z = estimated.at<float>(2);
  // 重置丢失计数
  this->missing_count = 0;
  this->state = estimated.clone();
  return statePt;
}

cv::Point3f Predictor::getState() const {
  if (!this->init) {
    throw std::runtime_error("Kalman Filter not initialized!");
  }

  cv::Point3f statePt;
  statePt.x = this->state.at<float>(0);
  statePt.y = this->state.at<float>(1);
  statePt.z = this->state.at<float>(2);
  return statePt;
}

int Predictor::getMissingCount() const {
  return this->missing_count;
}
